import os
import zlib
import marshal
import base64
import hashlib
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad

class ViewerEncryptor:
    def __init__(self):
        company = "RockChips"
        project = "RMSLViewer"
        version = "2025.09"
        secret = "AnimateX#2025!"

        # 组合成密钥源
        key_source = f"{company}_{project}_{version}_{secret}"

        # 生成密钥和IV
        self.key = hashlib.sha256(key_source.encode()).digest()  # 32字节
        self.iv = hashlib.md5(f"{secret}_{version}".encode()).digest()  # 16字节

    def encrypt_python_file(self, filename):
        """加密Python文件并返回加密数据"""
        print(f"正在加密 {filename}...")

        # 1. 读取源代码
        with open(filename, 'r', encoding='utf-8') as f:
            source_code = f.read()

        # 2. 编译为字节码
        code_object = compile(source_code, filename, 'exec')
        bytecode = marshal.dumps(code_object)

        # 3. 压缩
        compressed = zlib.compress(bytecode, 9)

        # 4. AES加密
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        encrypted = cipher.encrypt(pad(compressed, AES.block_size))

        # 5. Base64编码
        encoded = base64.b64encode(encrypted)

        # 6. 保存加密文件（可选）
        with open(f'{filename}.enc', 'wb') as f:
            f.write(encoded)

        print(f"加密完成: {filename}.enc")
        return encoded

    def create_embedded_loader(self, encrypted_data):
        """创建包含嵌入加密数据的加载器"""
        # 将密钥和IV编码
        key_hex = self.key.hex()
        iv_hex = self.iv.hex()

        # 将加密数据分成多行，避免单行过长
        data_str = encrypted_data.decode('ascii')
        chunk_size = 80  # 每行80个字符
        data_chunks = [data_str[i:i+chunk_size] for i in range(0, len(data_str), chunk_size)]
        data_lines = '\n    '.join([f'"{chunk}"' for chunk in data_chunks])

        loader_code = f'''
# Encrypted viewer module with embedded data
def _init_viewer():
    import marshal, zlib, base64, types, sys
    from Crypto.Cipher import AES
    from Crypto.Util.Padding import unpad

    # Embedded encrypted data
    _enc_data = (
    {data_lines}
    )

    # Reconstruct key and IV
    _k = bytes.fromhex('{key_hex}')
    _i = bytes.fromhex('{iv_hex}')

    # Decrypt and load
    _data = base64.b64decode(_enc_data)
    _cipher = AES.new(_k, AES.MODE_CBC, _i)
    _decrypted = unpad(_cipher.decrypt(_data), AES.block_size)
    _code = marshal.loads(zlib.decompress(_decrypted))

    # Create module
    viewer = types.ModuleType('viewer')
    viewer.__file__ = '<encrypted>'
    exec(_code, viewer.__dict__)

    # Register in sys.modules
    sys.modules['viewer'] = viewer
    return viewer

# Initialize viewer module
viewer = _init_viewer()

# Clean up temporary function
del _init_viewer
'''
        return loader_code

    def create_compact_embedded_loader(self, encrypted_data):
        """创建更紧凑的嵌入式加载器（用于pystand.int）"""
        # 使用更短的变量名和更紧凑的格式
        key_hex = self.key.hex()
        iv_hex = self.iv.hex()

        # 将数据压缩成一个长字符串
        data_str = encrypted_data.decode('ascii')

        # 创建极简版本
        loader_code = f'''
# Viewer module loader
import marshal,zlib,base64,types,sys
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
_d="{data_str}"
_c=AES.new(bytes.fromhex('{key_hex}'),AES.MODE_CBC,bytes.fromhex('{iv_hex}'))
_m=types.ModuleType('viewer')
exec(marshal.loads(zlib.decompress(unpad(_c.decrypt(base64.b64decode(_d)),16))),_m.__dict__)
sys.modules['viewer']=viewer=_m
'''
        return loader_code

    def create_modular_embedded_loader(self, encrypted_data):
        """创建模块化的嵌入式加载器（更易维护）"""
        # 将密钥分散存储
        key_parts = [self.key[i:i+8].hex() for i in range(0, len(self.key), 8)]
        iv_parts = [self.iv[i:i+8].hex() for i in range(0, len(self.iv), 8)]

        # 将数据也分块
        data_str = encrypted_data.decode('ascii')
        chunk_size = 60
        data_chunks = [data_str[i:i+chunk_size] for i in range(0, len(data_str), chunk_size)]

        loader_code = f'''
# Encrypted viewer module
def _load_viewer():
    import marshal, zlib, base64, types, sys
    from Crypto.Cipher import AES
    from Crypto.Util.Padding import unpad

    # Reconstruct data
    _p = []
    {chr(10).join([f'    _p.append("{chunk}")' for chunk in data_chunks])}
    _data = "".join(_p)

    # Reconstruct key
    _k = ""
    {chr(10).join([f'    _k += "{part}"' for part in key_parts])}

    # Reconstruct IV
    _i = ""
    {chr(10).join([f'    _i += "{part}"' for part in iv_parts])}

    # Decrypt and load
    cipher = AES.new(bytes.fromhex(_k), AES.MODE_CBC, bytes.fromhex(_i))
    decrypted = unpad(cipher.decrypt(base64.b64decode(_data)), 16)
    code = marshal.loads(zlib.decompress(decrypted))

    # Create and register module
    m = types.ModuleType('viewer')
    m.__file__ = '<encrypted>'
    exec(code, m.__dict__)
    sys.modules['viewer'] = m

    return m

# Load viewer
viewer = _load_viewer()
del _load_viewer
'''
        return loader_code

    def create_secure_package(self):
        """创建完整的加密包"""
        # 1. 加密viewer.py并获取加密数据
        encrypted_data = self.encrypt_python_file('viewer_rc.py')

        # 2. 创建标准嵌入式加载器
        embedded_loader = self.create_embedded_loader(encrypted_data)
        with open('viewer_embedded.py', 'w', encoding='utf-8') as f:
            f.write(embedded_loader)

        # 3. 创建紧凑版本（用于pystand.int）
        compact_loader = self.create_compact_embedded_loader(encrypted_data)
        with open('viewer_compact.py', 'w', encoding='utf-8') as f:
            f.write(compact_loader)

        # 4. 创建模块化版本（更安全）
        modular_loader = self.create_modular_embedded_loader(encrypted_data)
        with open('viewer_modular.py', 'w', encoding='utf-8') as f:
            f.write(modular_loader)

        print("\n加密包创建完成！")
        print("- viewer_rc.py.enc: 独立的加密文件")
        print("- viewer_embedded.py: 标准嵌入式加载器")
        print("- viewer_compact.py: 极简版本（适合嵌入pystand.int）")
        print("- viewer_modular.py: 模块化版本（密钥和数据分散存储）")
        print("\n使用方法：")
        print("1. 将viewer_compact.py的内容直接复制到pystand.int中")
        print("2. 或者在pystand.int中添加: exec(open('viewer_embedded.py').read())")
        print("3. 之后就可以直接 import viewer 使用了")

if __name__ == "__main__":
    encryptor = ViewerEncryptor()
    encryptor.create_secure_package()
